package com.gorkemozbek.rubic.util

import scala.collection.{Iterator, Set, Seq}

/*
 * A generic counter where elements can be mapped to non-negative integer
 * values.
 */
trait Counter[E] {

  protected def nonNegativeInt(x: Int) = Math max(0, x)

  protected def updateCount(elem: E, cnt: Int): Unit

  final def setCount(elem: E, cnt: Int): Unit = updateCount(elem, nonNegativeInt(cnt))

  def counts: Iterator[(E, Int)]

  def remove(elem: E): Boolean

  def clear: Unit = counts foreach { cnt => remove(cnt._1) }

  def getCount(elem: E): Int = counts find { cnt => cnt._1 == elem } match { case Some(x) => x._2; case None => 0 }

  def elements: Set[E] = Set() ++ counts map { cnt => cnt._1 }

  def increment(elem: E, inc: Int = 1): Unit = setCount(elem, getCount(elem) + inc)

  def incrementAll(inc: Int = 1): Unit = counts foreach { cnt => increment(cnt._1, inc) }

  def decrement(elem: E, dec: Int = 1): Unit = setCount(elem, getCount(elem) - dec)

  def decrementAll(dec: Int = 1): Unit = counts foreach { cnt => decrement(cnt._1, dec) }

  def contains(elem: E): Boolean = counts exists { cnt => cnt._1 == elem }

  def size: Int = counts.length

  def isEmpty: Boolean = ! counts.hasNext

  def totalCount: Long = counts.foldLeft(0L) { (totalCountSoFar: Long, cnt: (E, Int)) => totalCountSoFar + cnt._2.toLong }

  def averageCount: Double = totalCount.toDouble / size.toDouble

  protected def leastCount(lessThan: (Int, Int) => Boolean): (E, Int) = counts.reduceLeft { (cnt1: (E, Int), cnt2: (E, Int)) => if (lessThan(cnt1._2, cnt2._2)) cnt1 else cnt2 }

  protected def maxCount: (E, Int) = leastCount{ _ > _ }

  protected def minCount: (E, Int) = leastCount{ _ < _ }

  def max: Int = maxCount._2

  def min: Int = minCount._2

  def argmax: E = if (isEmpty) throw new NoSuchElementException else maxCount._1

  def argmin: E = if (isEmpty) throw new NoSuchElementException else minCount._1

  def sortCounts: Seq[(E, Int)] = sortCounts { _ < _ }

  def sortCounts(lessThan: (Int, Int) => Boolean): Seq[(E, Int)] = counts.toSeq sortWith { (cnt1: (E, Int), cnt2: (E, Int)) => lessThan(cnt1._2, cnt2._2) }

  def filter(test: (E, Int) => Boolean): Unit = counts foreach { (cnt: (E, Int)) => if (! test(cnt._1, cnt._2)) remove(cnt._1) }

  def choose(test: (E, Int) => Boolean): Set[(E, Int)] = Set() ++ counts filter { cnt => test(cnt._1, cnt._2) }

  def chooseLessThan(limit: Int): Set[(E, Int)] = choose { (_, c: Int) => c < limit }

  def chooseLessThanOrEqualTo(limit: Int): Set[(E, Int)] = choose { (_, c: Int) => c <= limit }

  def chooseGreaterThan(limit: Int): Set[(E, Int)] = choose { (_, c: Int) => c > limit }

  def chooseGreaterThanOrEqualTo(limit: Int): Set[(E, Int)] = choose { (_, c: Int) => c >= limit }

}
